import json
import os
import shutil
from functools import partial
from glob import glob
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import numpy as np
from skimage.metrics import structural_similarity
from tqdm import tqdm

from option import parse_args

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)

cache = {}


def save_diffs(pair, root_dir):
    ori_id, fake_id = pair
    ori_dir = os.path.join(root_dir, "crops", ori_id)
    fake_dir = os.path.join(root_dir, "crops", fake_id)
    diff_dir = os.path.join(root_dir, "diffs", fake_id)
    os.makedirs(diff_dir, exist_ok=True)
    for frame in range(320):
        if frame % 10 != 0:
            continue
        for actor in range(2):
            image_id = "{}_{}.png".format(frame, actor)
            diff_image_id = "{}_{}_diff.png".format(frame, actor)
            ori_path = os.path.join(ori_dir, image_id)
            fake_path = os.path.join(fake_dir, image_id)
            diff_path = os.path.join(diff_dir, diff_image_id)
            if os.path.exists(ori_path) and os.path.exists(fake_path):
                img1 = cv2.imread(ori_path, cv2.IMREAD_COLOR)
                img2 = cv2.imread(fake_path, cv2.IMREAD_COLOR)
                try:
                    d, a = structural_similarity(img1, img2, multichannel=True, full=True)
                    a = 1 - a
                    diff = (a * 255).astype(np.uint8)
                    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
                    cv2.imwrite(diff_path, diff)
                except:
                    pass


def copy_diffs(pair, root_dir):
    part_id, ori_id, fake_id = pair
    ori_dir = os.path.join(root_dir, "crops", part_id, ori_id)
    fake_dir = os.path.join(root_dir, "crops", part_id, fake_id)
    diff_dir = os.path.join(root_dir, "diffs", fake_id)
    part_diff_dir = os.path.join(root_dir, "diffs", part_id, fake_id)

    if os.path.exists(diff_dir):
        shutil.move(diff_dir, part_diff_dir)


def get_original_with_fakes(root_dir):
    videos_dir = os.path.join(root_dir, 'videos')
    pairs = []
    for json_path in glob(os.path.join(videos_dir, "*/metadata.json")):
        part_name = Path(json_path).parent.name
        with open(json_path, "r") as f:
            metadata = json.load(f)
        for k, v in metadata.items():
            original = v.get("original", None)
            if v["label"] == "FAKE":
                pairs.append((part_name, original[:-4], k[:-4]))
    return pairs


def main():
    args = parse_args()
    pairs = get_original_with_fakes(args.root_dir)
    os.makedirs(os.path.join(args.root_dir, "diffs"), exist_ok=True)
    with Pool(processes=1) as p:
        with tqdm(total=len(pairs)) as pbar:
            func = partial(copy_diffs, root_dir=args.root_dir)
            for v in p.imap_unordered(func, pairs):
                pbar.update()


if __name__ == '__main__':
    main()
